import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
import seaborn as sns
from python_ai.common.xcommon import sep

pd.set_option('display.max_rows', None, 'display.max_columns', None, 'display.max_colwidth', 1000, 'display.expand_frame_repr', False)
plt.rcParams['font.sans-serif'] = ['Simhei']
plt.rcParams['axes.unicode_minus'] = False

class NaiveBayesImplWithSmoothReview02(object):

    def __init__(self, lam=1):
        self.lam_ = lam

    def fit(self, df):
        m = len(df)
        x = df.iloc[:, :-1]
        y = df.iloc[:, -1]
        P = {}
        self.x_names_ = x.columns

        y_cnt_series = y.value_counts()
        y_v_s = y_cnt_series.index
        self.y_v_s_ = np.array(y_v_s)
        y_v_len = len(y_v_s)
        self.y_v_len_ = y_v_len
        y_p_series = (y_cnt_series + self.lam_) / (m + y_v_len * self.lam_)
        self.y_p_series_ = y_p_series

        for y_v in y_v_s:
            df4y_v = df[df.iloc[:, -1] == y_v]
            df4y_v_len = len(df4y_v)

            for x_name in x.columns:
                x_series = df4y_v[x_name]
                x_cnt_series = x_series.value_counts()
                x_v_len = len(x_cnt_series)
                k = (x_name, None, y_v)
                P[k] = self.lam_ / (df4y_v_len + x_v_len * self.lam_)
                for i, x_v in enumerate(x_cnt_series.index):
                    x_cnt = x_cnt_series.iloc[i]
                    k = (x_name, x_v, y_v)
                    P[k] = (x_cnt + self.lam_) / (df4y_v_len + x_v_len * self.lam_)
        self.P_ = P
        print(self.y_p_series_)  # tmp
        print(self.P_)  # tmp

    def get_P_from_key(self, k):
        if k in self.P_:
            return self.P_[k]
        else:
            x_name, x_v, y_v = k
            return self.P_[(x_name, None, y_v)]

    def predict(self, x):
        m = len(x)
        h = np.zeros(m)
        for ii, sample in enumerate(x):
            ps = np.zeros(self.y_v_len_)
            for i, y_v in enumerate(self.y_v_s_):
                p = self.y_p_series_[y_v]
                # print(f'p:{p}')
                for j, x_v in enumerate(sample):
                    x_name = self.x_names_[j]
                    k = x_name, x_v, y_v
                    v = self.get_P_from_key(k)
                    # print(v)
                    p *= v
                    # print(f'p[after_j]:{p}')
                ps[i] = p
            print(ps)  # tmp
            idx_sorted = ps.argsort()[-1]
            h[ii] = self.y_v_s_[idx_sorted]
        return h


if '__main__' == __name__:
    df = pd.read_csv('../follow_teacher/bayes_lihang.txt', header=0)
    print(len(df))
    print(df[:5])

    model = NaiveBayesImplWithSmoothReview02(lam=1)
    model.fit(df)

    Xs = [[2, 'S'], [2, 'N'], [200, 'N'], [2, 'L']]
    rs = model.predict(Xs)
    print(rs)


